跳转至

使用 Hugging Face Bucket 传输万亿参数：TRL 中的 Delta 权重同步机制

文章背景与核心概要

在异步强化学习（Async RL）中，训练引擎与推理引擎之间的权重同步一直是一个主要的性能瓶颈。传统方法要求训练器在每一步都向推理引擎传输完整的模型检查点（从 7B 模型的 14 GB 到万亿参数前沿模型的约 1 TB 不等）。然而，在连续的 RL 优化步骤之间，大约 99% 的 bf16 权重保持完全一致，实际的增量微乎其微。

为了解决这一问题，TRL 中的一项新功能将发生变化的元素编码为稀疏的 safetensors 文件，将其上传到 Hugging Face Bucket，并引导 vLLM 进行拉取。这大幅减少了单步传输的数据量（例如在 Qwen3-0.6B 上从 1.2 GB 降至 20–35 MB），并实现了跨不同区域的完全解耦训练架构，彻底摆脱了对共享集群、RDMA 网络或 VPN 的依赖。

1. 一万亿参数的难题 (The One Terabyte Problem)

Every async RL library eventually faces the same fundamental challenge: weight synchronization.

每一个异步强化学习（Async RL）库最终都会面临一个共同的根本挑战：权重同步。

The inference engine operates on the policy of step \(N\), while the trainer has just completed step \(N+1\). To prevent the inference engine from drifting off-policy, fresh weights must be transferred immediately. In traditional setups, a blocking transfer wastes idle compute cycles while GPUs wait rather than generating tokens.

推理引擎基于步骤 \(N\) 的策略运行，而训练器刚刚完成了步骤 \(N+1\)。为了防止推理引擎偏离策略（off-policy），必须立即传输最新的权重。在传统设置中，阻塞式传输会白白浪费空闲的计算周期，导致 GPU 忙于等待传输而非生成 Token。

For a frontier 1T-parameter checkpoint at fp8 precision, a full snapshot is roughly 1024 GiB. Conventional workflows demand shipping this massive payload on every update. Yet, measured data shows that the average delta between adjacent checkpoints is only 20.3 GiB (1.98% of the full model), with more than 98% of weights remaining bit-equivalent.

对于采用 fp8 精度的万亿参数前沿模型检查点，完整的快照大约为 1024 GiB。传统的工作流要求在每次更新时传输这个庞大的数据负载。然而，实测数据显示，相邻检查点之间的平均增量仅为 20.3 GiB（占完整模型的 1.98%），并且超过 98% 的权重在比特层面保持一致。

Leveraging shared object storage (like S3 or HF Buckets) allows trainers to upload compressed weight diffs asynchronously. The trainer simply publishes "weights ready" and uploads the diff the moment an optimizer step finishes, while inference replicas download and reconstruct updates independently on their own time.

利用共享对象存储（如 S3 或 HF Buckets），训练器可以异步上传压缩后的权重差异（diff）。训练器只需在优化器步骤完成的瞬间发布“权重就绪”消息并上传差分，而推理副本则可以在它们自己的时间里独立下载并重建更新。

2. 为什么 bf16 RL 权重几乎总是稀疏的 (Why bf16 RL Weights Are Almost Always Sparse)

This high degree of sparsity is rooted in how bf16 floating-point arithmetic interacts with the learning rates typically used in RL.

这种高度的稀疏性植根于 bf16 浮点运算与 RL 中通常使用的学习率之间的相互作用。

A bf16 number features 7 mantissa bits, meaning there are \(2^7 = 128\) representable values between any two consecutive powers of two. An update is absorbed by the bf16 cast whenever it sits below half of that spacing—specifically when \(|\Delta w| < |w|/256\) (the bf16 visibility threshold).

bf16 数字具有 7 个尾数位，这意味着在任意两个连续的 2 的幂之间有 \(2^7 = 128\) 个可表示的值。当更新值小于该间距的一半时，它就会被 bf16 类型转换所吸收——具体而言，即当 \(|\Delta w| < |w|/256\) 时（即 bf16 可见性阈值）。

At standard RL learning rates (\(\eta \approx 3 \times 10^{-6}\)), the update to a weight is given by: $\(\Delta w = -\eta \cdot \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon}\)$

在标准的 RL 学习率（\(\eta \approx 3 \times 10^{-6}\)）下，对权重的更新公式为： $\(\Delta w = -\eta \cdot \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon}\)$

Because the normalized step magnitude is roughly order-one, \(|\Delta w| \approx \eta \approx 3 \times 10^{-6}\). Given that typical LLM weights have magnitudes around \(10^{-2}\) to \(10^{-1}\), the corresponding visibility threshold (\(|w|/256\)) evaluates to \(4 \times 10^{-5}\) to \(4 \times 10^{-4}\)—which is larger than the update itself.

由于归一化步长量级大致为 1，因此 \(|\Delta w| \approx \eta \approx 3 \times 10^{-6}\)。鉴于典型 LLM 权重的量级在 \(10^{-2}\) 到 \(10^{-1}\) 之间，相应的可见性阈值（\(|w|/256\)）计算结果为 \(4 \times 10^{-5}\) 到 \(4 \times 10^{-4}\)——这大于更新值本身。

As a result, the optimizer's updates are swallowed by rounding errors. The byte representation of \(w\) remains unchanged for over 99% of the parameters, yielding near-total sparsity naturally and without approximations.

结果就是，优化器的更新被舍入误差所吞没。超过 99% 的参数的 \(w\) 字节表示保持不变，从而在没有进行任何近似的情况下，自然地产生了接近完全的稀疏性。

3. HF Buckets 与架构 (HF Buckets and the Architecture)

3.1 什么是 Bucket？ (3.1 What is a Bucket?)

A Bucket is a specialized Hub repository type designed for high-frequency object storage without standard commit ceremonies or LFS quirks. It uses Xet (content-defined chunking) under the hood to automatically slice uploaded files based on content and deduplicate them against existing objects. Even without sparse encoding, Xet ensures that only modified chunks travel across the wire.

Bucket 是一种专为高频对象存储设计的特殊 Hub 仓库类型，它不需要标准的提交仪式或 LFS 的怪异行径。它底层使用 Xet（基于内容的分块技术），根据内容自动切片上传的文件，并针对现有对象进行去重。即使没有稀疏编码，Xet 也能确保只有被修改的块在网络上传输。

3.2 三方箱体设置 (3.2 The Three-Box Setup)

Disaggregated training relies on three distinct components communicating through a shared bucket:

解耦训练依赖于通过共享存储桶进行通信的三个不同组件：

  1. Trainer: Owns the model, runs the optimizer, and emits sparse deltas.
    1. 训练器（Trainer）： 拥有模型、运行优化器并输出稀疏增量。
  2. HF Bucket: Stores occasional full snapshots (anchors/) and sparse patches (deltas/).
    1. HF Bucket： 存储不定期的完整快照（anchors/）和稀疏补丁（deltas/）。
  3. vLLM Rollout Server: Pulls updates from the bucket, applies deltas, and serves rollouts (can run anywhere, including inside a Hugging Face Space).
    1. vLLM Rollout 服务器： 从存储桶拉取更新、应用增量并提供 rollout 服务（可以运行在任何地方，包括 Hugging Face Space 内部）。

4. 协议 (The Protocol)

4.1 Safetensors 作为传输格式 (4.1 Safetensors as the Wire Format)

The system uses safetensors for both on-disk storage and network transport.

该系统将 safetensors 用于磁盘存储和网络传输。

  • Anchors: Full model snapshots written every \(N\) steps (default \(N=10\)).
    • 锚点（Anchors）： 每 \(N\) 步写入一次完整的模型快照（默认 \(N=10\)）。
  • Deltas: Sparse patches containing a flat int32 tensor of modified element indices and a bf16 tensor of updated values.
    • 增量（Deltas）： 包含被修改元素索引的扁平 int32 张量以及更新后值的 bf16 张量的稀疏补丁。

4.2 训练器侧的变更检测器 (4.2 The Trainer-Side Change Detector)

A lightweight BF16ChangeDetector hooks into the optimizer to snapshot weights pre- and post-step, generating an exact boolean change mask:

一个轻量级的 BF16ChangeDetector 挂载到优化器中，用于在步骤执行前后对权重进行快照，从而生成精确的布尔值变更掩码：

class BF16ChangeDetector:
    def __init__(self, model, optimizer):
        self._pre_step_bf16: dict[str, torch.Tensor] = {}
        self._validated_masks: dict[str, torch.Tensor] = {}
        optimizer.register_step_pre_hook(self._pre_step_hook)
        optimizer.register_step_post_hook(self._post_step_hook)

    def _pre_step_hook(self, opt, args, kwargs):
        for p in self._params:
            self._pre_step_bf16[name_of(p)] = p.detach().to(torch.bfloat16).cpu().clone()

    def _post_step_hook(self, opt, args, kwargs):
        for p in self._params:
            self._validated_masks[name_of(p)] = (
                p.detach().to(torch.bfloat16).cpu() != self._pre_step_bf16[name_of(p)]
            )

4.3 vLLM 侧的 Worker 扩展 (4.3 The vLLM-Side Worker Extension)

Using vLLM's --worker-extension-cls flag, a custom extension downloads sparse patches from the bucket, reconstructs full tensors via a local CPU snapshot, and reloads them seamlessly into memory without requiring a custom vLLM fork.

通过使用 vLLM 的 --worker-extension-cls 标志，自定义扩展可以从存储桶下载稀疏补丁，通过本地 CPU 快照重建完整的张量，并无缝地将其重新加载到内存中，而无需对 vLLM 进行自定义分支修改。

5. 在 Spaces 上真正运行起来 (Standing It Up on Spaces, For Real)

Disaggregated training can be deployed across completely isolated environments without shared networking:

解耦训练可以部署在完全隔离、无共享网络的环境中：

  • Trainer: Running locally or on a dedicated GPU instance.
    • 训练器（Trainer）： 运行在本地或专用的 GPU 实例上。
  • vLLM Inference Space: A Hugging Face Space running an L4 GPU container with the TRL delta extension.
    • vLLM 推理 Space： 运行带有 TRL 增量扩展的 L4 GPU 容器的 Hugging Face Space。
  • Environment Space: A separate CPU Space handling interactive rollouts (e.g., Wordle environments).
    • 环境 Space： 处理交互式 rollout（例如 Wordle 环境）的独立 CPU Space。

To launch the inference space via the CLI:

通过 CLI 启动推理 Space：

hf repos create $USER/vllm-wordle-inference \
    --type space --space-sdk docker --flavor l4x1 \
    --secrets HF_TOKEN=$HF_TOKEN
hf upload $USER/vllm-wordle-inference examples/scripts/openenv/vllm_space/ --type space

And start the trainer:

并启动训练器：

python examples/scripts/openenv/async_wordle.py \
    --vllm-server-url https://$USER-vllm-wordle-inference.hf.space \
    --env-url https://openenv-wordle.hf.space \
    --delta-sync-repo-id $USER/wordle-deltas \
    --model Qwen/Qwen3-1.7B

6. 这带来了哪些突破 (What This Unlocks)

  • Cluster-Free Async RL: Run distributed RL setups using standard cloud hardware, HF Spaces, and object storage instead of expensive RDMA fabrics.
    • 无集群异步 RL： 使用标准云硬件、HF Spaces 和对象存储运行分布式 RL 设置，而无需昂贵的 RDMA 网络。
  • Multi-Replica Inference: Scale out multiple rollout replicas that pull from the same bucket simultaneously while benefiting from edge caching and deduplication.
    • 多副本推理： 扩展多个 rollout 副本，它们可以同时从同一个存储桶拉取数据，同时受益于边缘缓存和去重功能。
  • Transparent Debugging: Inspect protocol payloads directly via safetensors APIs without struggling through opaque NCCL network streams.
    • 透明化调试： 直接通过 safetensors API 检查协议负载，而无需在晦涩难懂的 NCCL 网络流中苦苦挣扎。